import pandas as pd
import numpy as np
from numpy import dtype, nan
import pytest
from pvlib.iotools import crn
from tests.conftest import (
    TESTS_DATA_DIR,
    assert_frame_equal,
    assert_index_equal,
)


@pytest.fixture
def columns_mapped():
    return [
        'WBANNO', 'UTC_DATE', 'UTC_TIME', 'LST_DATE', 'LST_TIME', 'CRX_VN',
        'longitude', 'latitude', 'temp_air', 'PRECIPITATION', 'ghi',
        'ghi_flag',
        'SURFACE_TEMPERATURE', 'ST_TYPE', 'ST_FLAG', 'relative_humidity',
        'relative_humidity_flag', 'SOIL_MOISTURE_5', 'SOIL_TEMPERATURE_5',
        'WETNESS', 'WET_FLAG', 'wind_speed', 'wind_speed_flag']


@pytest.fixture
def columns_unmapped():
    return [
        'WBANNO', 'UTC_DATE', 'UTC_TIME', 'LST_DATE', 'LST_TIME', 'CRX_VN',
        'LONGITUDE', 'LATITUDE', 'AIR_TEMPERATURE', 'PRECIPITATION',
        'SOLAR_RADIATION', 'SR_FLAG', 'SURFACE_TEMPERATURE', 'ST_TYPE',
        'ST_FLAG', 'RELATIVE_HUMIDITY', 'RH_FLAG', 'SOIL_MOISTURE_5',
        'SOIL_TEMPERATURE_5', 'WETNESS', 'WET_FLAG', 'WIND_1_5', 'WIND_FLAG']


@pytest.fixture
def dtypes():
    return [
        dtype('int64'), dtype('int64'), dtype('int64'), dtype('int64'),
        dtype('int64'), dtype('O'), dtype('float64'), dtype('float64'),
        dtype('float64'), dtype('float64'), dtype('float64'),
        dtype('int64'), dtype('float64'), dtype('O'), dtype('int64'),
        dtype('float64'), dtype('int64'), dtype('float64'),
        dtype('float64'), dtype('int64'), dtype('int64'), dtype('float64'),
        dtype('int64')]


@pytest.fixture
def testfile():
    return TESTS_DATA_DIR / 'CRNS0101-05-2019-AZ_Tucson_11_W.txt'


@pytest.fixture
def testfile_problems():
    return TESTS_DATA_DIR / 'CRN_with_problems.txt'


def test_read_crn(testfile, columns_mapped, dtypes):
    index = pd.DatetimeIndex(['2019-01-01 16:10:00',
                              '2019-01-01 16:15:00',
                              '2019-01-01 16:20:00',
                              '2019-01-01 16:25:00'],
                             freq=None).tz_localize('UTC')
    values = np.array([
        [53131, 20190101, 1610, 20190101, 910, 3, -111.17, 32.24, nan,
         0.0, 296.0, 0, 4.4, 'C', 0, 90.0, 0, nan, nan, 24, 0, 0.78, 0],
        [53131, 20190101, 1615, 20190101, 915, 3, -111.17, 32.24, 3.3,
         0.0, 183.0, 0, 4.0, 'C', 0, 87.0, 0, nan, nan, 1182, 0, 0.36, 0],
        [53131, 20190101, 1620, 20190101, 920, 3, -111.17, 32.24, 3.5,
         0.0, 340.0, 0, 4.3, 'C', 0, 83.0, 0, nan, nan, 1183, 0, 0.53, 0],
        [53131, 20190101, 1625, 20190101, 925, 3, -111.17, 32.24, 4.0,
         0.0, 393.0, 0, 4.8, 'C', 0, 81.0, 0, nan, nan, 1223, 0, 0.64, 0]])
    expected = pd.DataFrame(values, columns=columns_mapped, index=index)
    for (col, _dtype) in zip(expected.columns, dtypes):
        expected[col] = expected[col].astype(_dtype)
    out = crn.read_crn(testfile)
    assert_frame_equal(out, expected)


# Test map_variables=False returns correct column names
def test_read_crn_map_variables(testfile, columns_unmapped, dtypes):
    out = crn.read_crn(testfile, map_variables=False)
    assert_index_equal(out.columns, pd.Index(columns_unmapped))


def test_read_crn_problems(testfile_problems, columns_mapped, dtypes):
    # GH1025
    index = pd.DatetimeIndex(['2020-07-06 12:00:00',
                              '2020-07-06 13:10:00'],
                             freq=None).tz_localize('UTC')
    values = np.array([
        [92821, 20200706, 1200, 20200706, 700, '3', -80.69, 28.62, 24.9,
         0.0, np.nan, 0, 25.5, 'C', 0, 93.0, 0, nan, nan, 990, 0, 1.57, 0],
        [92821, 20200706, 1310, 20200706, 810, '2.623', -80.69, 28.62,
         26.9, 0.0, 430.0, 0, 30.2, 'C', 0, 87.0, 0, nan, nan, 989, 0,
         1.64, 0]])
    expected = pd.DataFrame(values, columns=columns_mapped, index=index)
    for (col, _dtype) in zip(expected.columns, dtypes):
        expected[col] = expected[col].astype(_dtype)
    out = crn.read_crn(testfile_problems)
    assert_frame_equal(out, expected)
